import numpy as np
from scipy.spatial.distance import cdist
from scipy.optimize import minimize
import matplotlib.pyplot as plt
from scipy.optimize import linear_sum_assignment
from matplotlib.animation import FuncAnimation
'''
优化目标： 一条路径中最短距离
优化方法：
'''
# 定义全局参数
num_drones = 10
num_targets = 40
pheromone_decay = 0.1
pheromone_increase = 1.0
initial_pheromone = 1.0
alpha = 1.0  # 信息素重要程度
beta = 2.0  # 启发函数重要程度
Q = 1.0  # 信息素释放强度

plt.figure(figsize=(200, 100))  # 设置图的大小为100x100
def distance_objective(assignments, distance_matrix, type="normal"):
    """
    Objective function to minimize the difference in the sums of distances for each drone.
    """
    if type == "enemy":
        # assignments = assignments  # Ensure assignments is of integer type
        selected_distances = distance_matrix[np.arange(len(assignments)), assignments.astype(int)]
        return np.sum(selected_distances)
    drone_sums = np.array([distance_matrix[assignments.astype(int) == i].sum() for i in range(int(assignments.max()) + 1)])
    return np.std(drone_sums)

def custom_objective(assignments, distance_matrix, num_nomal):
    """
    Custom objective function for the constrained assignment.
    """
    normal_distances = distance_matrix[:num_nomal, :]
    enemy_distances = distance_matrix[num_nomal:, :]
    # Objective for enemy drones
    enemy_objective = np.sum(enemy_distances[np.arange(len(assignments[num_nomal:])), assignments[num_nomal:].astype(int)])

    # Objective for normal drones
    normal_objective = distance_objective(assignments[:num_nomal], normal_distances, type="normal")

    return enemy_objective + normal_objective

def find_closest_aircraft(enemy_distance_matrix):
    # 找到每个目标点对应的最近飞机下标
    closest_aircraft_indices = np.argsort(enemy_distance_matrix, axis=1)
    # 选择每列的第一个下标，确保不相同
    selected_indices = []
    for col_indices in closest_aircraft_indices:
        for index in col_indices:
            if index not in selected_indices:
                selected_indices.append(index)
                break

    # 检查是否有相同的选择下标，如果有则比较数字之和
    for i in range(len(closest_aircraft_indices)):
        for j in range(i + 1, len(closest_aircraft_indices)):
            if selected_indices[i] == selected_indices[j]:
                sum_i = np.sum(enemy_distance_matrix[:, i])
                sum_j = np.sum(enemy_distance_matrix[:, j])
                if sum_i < sum_j:
                    selected_indices[j] = closest_aircraft_indices[j, 1]
                else:
                    selected_indices[i] = closest_aircraft_indices[i, 1]
    return  np.array(selected_indices)

def get_new_drones(drones, indices_min_distance):
    # 记录新旧下标对应关系的字典
    index_mapping = {}
    # 新的 drones 数组
    new_drones = np.delete(drones, indices_min_distance, axis=0)
    # 构建新旧下标对应关系
    for i, old_index in enumerate(range(drones.shape[0])):
        if old_index not in indices_min_distance:
            new_index = np.where(np.delete(np.arange(drones.shape[0]), indices_min_distance) == old_index)[0][0]
            index_mapping[new_index] = old_index

    return new_drones, index_mapping
def assign_targets_to_drones(targets, drones, targets_enemy):
    # num_targets = targets.shape[0]
    # num_drones = drones.shape[0]
    new_drones = drones
    final_assignments_enemy = []
    # 计算距离矩阵
    #print("targets, drones",targets, drones)
    # distance_matrix = cdist(targets, drones)
    if len(targets_enemy) != 0:
        enemy_distance_matrix = cdist(targets_enemy, drones)
        # Use linear_sum_assignment to find the optimal assignment
        final_assignments_enemy = find_closest_aircraft(enemy_distance_matrix)
        new_drones, index_mapping = get_new_drones(drones, final_assignments_enemy)
        # print("distance_matrix", distance_matrix)
    distance_matrix = cdist(targets, new_drones)
    # 初始分配
    initial_assignments = np.argmin(distance_matrix, axis=1)
    # print("initial_assignments", initial_assignments)
    # 定义目标函数
    objective_function = lambda x: custom_objective(x, distance_matrix, len(targets))

    # 最小化目标函数
    result = minimize(objective_function, initial_assignments, method='L-BFGS-B')

    # 最终分配
    final_assignments = result.x.astype(int)
    if len(targets_enemy) != 0:
        # 使用映射关系字典获取原来的下标
        final_assignments = np.array([index_mapping[new_index] for new_index in final_assignments])
    return final_assignments, final_assignments_enemy
def draw_point(targets, drones, best_paths, time, iteration=1, name="Drone", targets_enemy=[], assignments_enemy=[]):
    plt.clf()
    for target in targets_enemy:
        plt.scatter(target[0], target[1], marker='*', color='red', s=100)
    # 绘制目标点
    for target in targets:
        plt.scatter(target[0], target[1], marker='x', color='red', s=100)
    # 绘制无人机起点
    for i in range(len(drones)):
        plt.scatter(drones[i][0], drones[i][1], marker='o', label=f'Drone {i + 1}', s=100)
    # 绘制路径
    for i, path in enumerate(best_paths):
        path_x, path_y = [drones[i][0]], [drones[i][1]]
        if path is not None:
            # if name == "Drone":
            #     drones[i] = targets[path[0]]
            for node in path:
                x, y = targets[node]
                path_x.append(x)
                path_y.append(y)
        plt.plot(path_x, path_y, linestyle='-', marker='o')
    for i, index in enumerate(assignments_enemy):
        path_x, path_y = [drones[index][0]], [drones[index][1]]
        if name == "Drone":
            # drones[index] = targets_enemy[i]
            x, y = targets_enemy[i]
            path_x.append(x)
            path_y.append(y)
        plt.plot(path_x, path_y, linestyle='--', marker='s')
    if name == "iteration":
        plt.title(f'Drone Paths - Iteration {iteration + 1}')
    else:
        plt.title(f'Drone Paths {iteration}')
    plt.xlabel('X-coordinate')
    plt.ylabel('Y-coordinate')
    plt.legend()
    plt.grid(True)
    plt.show(block=False)
    plt.pause(time)





# 计算两点之间的距离
def distance(point1, point2):
    return np.sqrt(np.sum((point1 - point2) ** 2))

# 更新信息素
def update_pheromones(paths, path_lengths, pheromones_target):
    # pheromones *= (1.0 - pheromone_decay)  # 信息素挥发
    pheromones_target *= (1.0 - pheromone_decay)

    # for j in range(len(num_targets)):
    #     pheromones_target[j] += pheromone_increase # 信息素更新方式，直接累加
    # 更新pheromones_target
    for j in range(len(paths)):
        i = paths[j]  # 获取paths中的数字
        pheromones_target[i] += pheromone_increase / path_lengths[j]
    return pheromones_target


# def update(frame):
def update(num_drones, targets, best_paths, iteration):
    plt.clf()

    # Plot targets
    for target in targets:
        plt.scatter(target[0], target[1], marker='x', color='red', s=100)

    # Plot drones' starting points
    for i in range(num_drones):
        plt.scatter(drones[i][0], drones[i][1], marker='o', label=f'Drone {i + 1}', s=100)

    # Plot paths
    for i, path in enumerate(best_paths):
        path_x, path_y = [drones[i % num_drones][0]], [drones[i % num_drones][1]]
        if path is not None:
            for node in path:
                x, y = targets[node]
                path_x.append(x)
                path_y.append(y)
            plt.plot(path_x, path_y, linestyle='-', marker='o')

    plt.title(f'Drone Paths - Iteration {iteration + 1}')
    plt.xlabel('X-coordinate')
    plt.ylabel('Y-coordinate')
    plt.legend()
    plt.grid(True)
    plt.show(block=False)
    plt.pause(0.01)

# 蚁群算法
def ant_colony( targets, drones, pheromones_target, targets_enemy=[]):
    num_iterations = 100
    # Before the loop
    best_paths = np.empty((len(drones)), dtype=object)
    best_path_lengths = np.ones((len(drones))) * np.inf
    assignments,  assignments_enemy = assign_targets_to_drones(targets, drones, targets_enemy)
    # if len(targets_enemy) != 0:
    #     assignments_enemy = assign_targets_to_drones(targets_enemy, drones, "enemy")
    # print("assignments", assignments)
    for iteration in range(num_iterations):
        for i in range(len(drones)):
            path = []
            visited = set()
            current_position = drones[i]
            # 创建长度为num_的零数组
            count_arr = np.zeros(len(drones))
            # 使用np.bincount得到0和1的个数，并将其赋值给零数组
            count_arr[:len(np.bincount(assignments))] = np.bincount(assignments)
            # 将结果转换为整数类型
            count_arr = count_arr.astype(int)
            if count_arr[i]:
                for _ in range(count_arr[i]):
                    visited.add(tuple(current_position))
                    unvisited_targets = [[number, t.tolist()] for number, t in enumerate(targets) if
                                         tuple(t) not in visited and assignments[number] == i]

                    # print("unvisited_targets", unvisited_targets)
                    unvisited_targets_id = [number for number, _ in unvisited_targets]
                    if len(unvisited_targets_id) ==0:
                        continue
                    # print("unvisited_targets_id", unvisited_targets_id)
                    # 计算每个可选目标的概率
                    probabilities = np.zeros(len(unvisited_targets))
                    for k, [num, _] in enumerate(unvisited_targets):
                        pheromone_value = pheromones_target[num]
                        distance_value = distance(current_position, targets[num])

                        if pheromone_value != 0 and distance_value != 0:
                            probabilities[k] = (pheromone_value ** alpha) / (distance_value ** beta)
                        else:
                            probabilities[k] = 0

                    # 选择下一个目标

                    selected_target = np.random.choice(unvisited_targets_id, p=probabilities / np.sum(probabilities))
                    path.append(selected_target)
                    current_position = targets[selected_target]

                # 计算路径长度
                path_length = distance(drones[i], path[0])
                t_old = path[0]
                for t in path[1:]:
                    path_length += distance(targets[t_old], targets[t])
                    t_old = t

                # 更新最佳路径
                if path_length < best_path_lengths[i]:
                    best_path_lengths[i] = path_length
                    best_paths[i] = np.array(path)
                    #draw_point(targets, drones, best_paths, 0.01, iteration=iteration, name="iteration", targets_enemy=targets_enemy)


        # 更新信息素
        pheromones_target = update_pheromones( best_paths, best_path_lengths, pheromones_target)


    return best_paths, best_path_lengths, assignments_enemy

def remove_xy(targets, drones, best_paths,targets_enemy, assignments_enemy):
    for i, path in enumerate(best_paths):
        if path is not None:
            direction_vector = targets[path[0]] - drones[i]
            normalized_direction = direction_vector / np.linalg.norm(direction_vector)
            displacement_vector = 5 * normalized_direction
            drones[i] += displacement_vector

    for i, index in enumerate(assignments_enemy):
            direction_vector = targets_enemy[i] - drones[index]
            # print("i, index", i, index)
            # print("direction_vector", direction_vector)
            # print("np.linalg.norm(direction_vector)", np.linalg.norm(direction_vector))
            if np.linalg.norm(direction_vector) !=0 :
                normalized_direction = direction_vector / np.linalg.norm(direction_vector)
            else:
                normalized_direction = 0
            displacement_vector = 5 * normalized_direction
            drones[index] += displacement_vector
    path_remove_list = []
    for i, path in enumerate(best_paths):
        if path is not None:
            if distance(targets[path[0]], drones[i]) < 10:
                path_remove_list.append(path[0])
    targets = [target for i, target in enumerate(targets) if i not in path_remove_list]
    return targets

def remove_xy(targets, drones, assignments_enemy):
    for i, path in enumerate(best_paths):
        if path is not None:
            direction_vector = targets[path[0]] - drones[i]
            normalized_direction = direction_vector / np.linalg.norm(direction_vector)
            displacement_vector = 5 * normalized_direction
            drones[i] += displacement_vector

    for i, index in enumerate(assignments_enemy):
            direction_vector = targets_enemy[i] - drones[index]
            # print("i, index", i, index)
            # print("direction_vector", direction_vector)
            # print("np.linalg.norm(direction_vector)", np.linalg.norm(direction_vector))
            if np.linalg.norm(direction_vector) !=0 :
                normalized_direction = direction_vector / np.linalg.norm(direction_vector)
            else:
                normalized_direction = 0
            displacement_vector = 5 * normalized_direction
            drones[index] += displacement_vector
    return targets

def move_towards_target_single(current_position, target_position, step_size, flag, random_factor=50):
    direction = target_position - current_position
    # print("target_position - current_position", target_position, current_position)
    distance = np.linalg.norm(direction)
    # print("distance", distance)
    if abs(distance) < step_size:
        return target_position
    else:
        unit_direction = direction / distance
        # 引入随机性，使用正态分布生成随机偏移
        random_offset = np.random.normal(0, random_factor, size=2)
        if flag =="enemy":
            new_position = current_position - step_size * unit_direction + random_offset
        else:
            new_position = current_position + step_size * unit_direction + random_offset
        # print("new_position", new_position)
        return new_position

# 随机生成目标点和无人机初始位置
def move_towards_target(current, targets_point, initial_assignments, step_size, flag="nomal", random_factor=50):
    for i in range(len(current)):
        target_index = initial_assignments[i]
        target_position = targets_point[target_index]
        # print("current[i]1", current[i])
        # print("targets_point", targets_point)
        # print("target_position1", target_position)
        current[i] = move_towards_target_single(current[i], target_position, step_size, flag, random_factor)
        # print("current[i]", i)
        # print("current[i]2", current[i])
    return current


targets = np.random.rand(num_targets, 2) * 100
targets_copy = targets
drones = np.random.rand(num_drones, 2) * 100
targets_enemy = np.random.rand(4, 2) * 100 - 100
Step = 0
for index_ep in range(10):
    # targets_enemy = []
    print("index_ep+++++++++++++++", index_ep)
    if index_ep >=0:
        # print("targets", targets)
        # print("targets_enemy", targets_enemy)
        distance_matrix = cdist(targets_copy, targets_enemy)
        initial_assignments = np.argmin(distance_matrix, axis=1)
        # print("distance_matrix", distance_matrix)
        # print("initial_assignments", initial_assignments)
        targets = move_towards_target(targets_copy, targets_enemy, initial_assignments, 20)
        targets_copy = targets
        # print("targets", targets)
        Step = 0

    while(len(targets)):
        Step += 1
        if Step% 5 ==0:
            # print(len(drones), len(targets))
            # pheromones = np.ones((len(drones), len(targets))) * initial_pheromone
            if index_ep >= 0:
                #print("targets_enemy", targets_enemy)
                #print("drones", drones)
                distance_matrix = cdist(targets_enemy, drones)
                print("distance_matrix", distance_matrix)
                initial_assignments = np.argmin(distance_matrix, axis=1)
                #print("initial_assignments", initial_assignments)
                targets_enemy = move_towards_target(targets_enemy, drones, initial_assignments, 5)
                #print("targets_enemy", targets_enemy)
            pheromones_target = np.ones(len(targets)) * initial_pheromone
            print("Step:", Step)
            # print("tragets:", targets)
            # print("drones:", drones)
            # 运行蚁群算法
            best_paths, best_path_lengths, assignments_enemy = ant_colony( targets, drones, pheromones_target, targets_enemy=targets_enemy)
            # print("best_paths:", best_paths)
            # for i in range(len(best_paths)):
            #     print(f"Drone {i }:", best_paths[i])
            draw_point(targets, drones, best_paths, 0.05, iteration=index_ep, targets_enemy=targets_enemy, assignments_enemy=assignments_enemy)

            targets = remove_xy(targets, drones, best_paths, targets_enemy, assignments_enemy)
        # else:
        #     draw_point(targets, drones, best_paths, 0.05, iteration=index_ep, targets_enemy=targets_enemy,assignments_enemy=assignments_enemy)